#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
WaveSpeedAI API 处理器
用于与 WaveSpeedAI API 交互，包括文件上传和视频生成
"""

import os
import requests
import json
import time
from pathlib import Path
from typing import Optional, Dict, Any

# 加载环境变量
from dotenv import load_dotenv
load_dotenv()


class WaveSpeedAI:
    """
    WaveSpeedAI API 处理器类
    提供文件上传和视频生成功能
    """
    
    def __init__(self, api_key: Optional[str] = None):
        """
        初始化 WaveSpeedAI 处理器
        
        Args:
            api_key (str, optional): API 密钥，如果未提供则从环境变量获取
        """
        self.api_key = api_key or os.getenv("WAVESPEED_API_KEY")
        if not self.api_key:
            raise ValueError("API key must be provided or set in WAVESPEED_API_KEY environment variable")
        
        self.base_url = "https://api.wavespeed.ai/api/v3"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}"
        }
    
    def upload_binary_file(self, file_path: str, api_url: str = None) -> Optional[Dict[Any, Any]]:
        """
        上传二进制文件到 API 端点
        
        Args:
            file_path (str): 要上传的文件路径
            api_url (str, optional): API 上传地址，默认为 WaveSpeedAI 上传端点
            
        Returns:
            dict: 上传结果，包含下载 URL 等信息，失败时返回 None
        """
        if not os.path.exists(file_path):
            print(f"Error: File {file_path} does not exist")
            return None
            
        if api_url is None:
            api_url = f"{self.base_url}/media/upload/binary"
            
        try:
            with open(file_path, 'rb') as f:
                files = {'file': f}
                response = requests.post(api_url, files=files, headers=self.headers)
                
            if response.status_code == 200:
                result = response.json()
                print("Upload successful! Download Url:", result)
                # 返回上传结果中的URL
                if "data" in result and "download_url" in result["data"]:
                    return result["data"]["download_url"]
                elif "url" in result:
                    return result["url"]
                else:
                    # 如果没有找到URL字段，返回整个结果
                    return result
            else:
                print(f"Upload failed with status code: {response.status_code}")
                print("Response:", response.text)
                return None
                
        except Exception as e:
            print(f"Error occurred: {str(e)}")
            return None
    


    def generate_lipsync_video_from_folder(self, folder_path: str, order: str = "meanwhile", 
                                          resolution: str = "480p", seed: int = -1) -> Optional[str]:
        """
        从生成的音频和图片文件夹直接生成唇形同步视频
        
        Args:
            folder_path (str): 包含音频和图片文件的文件夹路径
            order (str): 音频顺序，可选 "meanwhile", "left_right", "right_left"
            resolution (str): 视频分辨率，可选 "480p", "720p"
            seed (int): 随机种子，-1 表示随机
            
        Returns:
            str: 生成的视频 URL，失败时返回 None
        """
        folder = Path(folder_path)
        
        # 检查文件夹是否存在
        if not folder.exists():
            print(f"Error: Folder {folder_path} does not exist")
            return None
        
        # 检查 unique.mp3 文件是否存在（单人说话场景）
        unique_audio = folder / "unique.mp3"
        background_image = folder / "background.png"
        
        if unique_audio.exists() and background_image.exists():
            print("检测到单人说话场景")
            return self.generate_single_video(
                audio_path=str(unique_audio),
                image_path=str(background_image),
                resolution=resolution,
                seed=seed
            )
        
        # 检查对话文件是否存在（两人对话场景）
        left_audio = folder / "left.mp3"
        right_audio = folder / "right.mp3"
        background_image = folder / "background.png"
        
        if left_audio.exists() and right_audio.exists() and background_image.exists():
            print("检测到两人对话场景")
            return self.generate_conversation_video(
                left_audio_path=str(left_audio),
                right_audio_path=str(right_audio),
                image_path=str(background_image),
                order=order,
                resolution=resolution,
                seed=seed
            )
        
        print("错误: 未找到所需的音频和图片文件")
        print(f"文件夹内容: {list(folder.iterdir())}")
        return None
    
    def generate_single_video(self, audio_path: str, image_path: str, 
                             resolution: str = "480p", seed: int = -1) -> Optional[str]:
        """
        生成单人说话视频
        
        Args:
            audio_path (str): 音频文件路径
            image_path (str): 背景图片路径
            resolution (str): 视频分辨率
            seed (int): 随机种子
            
        Returns:
            str: 生成的视频 URL，失败时返回 None
        """
        print("生成单人说话视频")
        print(f"音频文件: {audio_path}")
        print(f"图片文件: {image_path}")
        
        # 上传音频文件
        print("正在上传音频文件...")
        audio_result = self.upload_binary_file(audio_path)
        if not audio_result:
            print("音频文件上传失败")
            return None
        
        # 处理上传结果，获取真实URL
        if isinstance(audio_result, str):
            audio_url = audio_result
        elif isinstance(audio_result, dict):
            audio_url = audio_result.get("url") or audio_result.get("download_url") or audio_result.get("data", {}).get("download_url")
        else:
            print("音频文件URL格式错误")
            return None
            
        if not audio_url:
            print("无法获取音频文件 URL")
            return None
        
        # 上传图片文件
        print("正在上传图片文件...")
        image_result = self.upload_binary_file(image_path)
        if not image_result:
            print("图片文件上传失败")
            return None
        
        # 处理上传结果，获取真实URL
        if isinstance(image_result, str):
            image_url = image_result
        elif isinstance(image_result, dict):
            image_url = image_result.get("url") or image_result.get("download_url") or image_result.get("data", {}).get("download_url")
        else:
            print("图片文件URL格式错误")
            return None
            
        if not image_url:
            print("无法获取图片文件 URL")
            return None
        
        # 调用单角色视频生成 API
        print("正在调用单角色视频生成 API...")
        return self.generate_video(
            audio_url=audio_url,
            image_url=image_url,
            resolution=resolution,
            seed=seed
        )
    
    def generate_conversation_video(self, left_audio_path: str, right_audio_path: str, 
                                   image_path: str, order: str = "meanwhile",
                                   resolution: str = "480p", seed: int = -1) -> Optional[str]:
        """
        生成两人对话视频
        
        Args:
            left_audio_path (str): 左侧说话人音频文件路径
            right_audio_path (str): 右侧说话人音频文件路径
            image_path (str): 背景图片路径
            order (str): 音频顺序，可选 "meanwhile", "left_right", "right_left"
            resolution (str): 视频分辨率
            seed (int): 随机种子
            
        Returns:
            str: 生成的视频 URL，失败时返回 None
        """
        print("生成两人对话视频")
        print(f"左侧音频文件: {left_audio_path}")
        print(f"右侧音频文件: {right_audio_path}")
        print(f"图片文件: {image_path}")
        print(f"音频顺序: {order}")
        
        # 上传左侧音频文件
        print("正在上传左侧音频文件...")
        left_audio_result = self.upload_binary_file(left_audio_path)
        if not left_audio_result:
            print("左侧音频文件上传失败")
            return None
        
        # 处理上传结果，获取真实URL
        if isinstance(left_audio_result, str):
            left_audio_url = left_audio_result
        elif isinstance(left_audio_result, dict):
            left_audio_url = left_audio_result.get("url") or left_audio_result.get("download_url") or left_audio_result.get("data", {}).get("download_url")
        else:
            print("左侧音频文件URL格式错误")
            return None
            
        if not left_audio_url:
            print("无法获取左侧音频文件 URL")
            return None
        
        # 上传右侧音频文件
        print("正在上传右侧音频文件...")
        right_audio_result = self.upload_binary_file(right_audio_path)
        if not right_audio_result:
            print("右侧音频文件上传失败")
            return None
        
        # 处理上传结果，获取真实URL
        if isinstance(right_audio_result, str):
            right_audio_url = right_audio_result
        elif isinstance(right_audio_result, dict):
            right_audio_url = right_audio_result.get("url") or right_audio_result.get("download_url") or right_audio_result.get("data", {}).get("download_url")
        else:
            print("右侧音频文件URL格式错误")
            return None
            
        if not right_audio_url:
            print("无法获取右侧音频文件 URL")
            return None
        
        # 上传图片文件
        print("正在上传图片文件...")
        image_result = self.upload_binary_file(image_path)
        if not image_result:
            print("图片文件上传失败")
            return None
        
        # 处理上传结果，获取真实URL
        if isinstance(image_result, str):
            image_url = image_result
        elif isinstance(image_result, dict):
            image_url = image_result.get("url") or image_result.get("download_url") or image_result.get("data", {}).get("download_url")
        else:
            print("图片文件URL格式错误")
            return None
            
        if not image_url:
            print("无法获取图片文件 URL")
            return None
        
        # 调用多角色视频生成 API
        print("正在调用多角色视频生成 API...")
        return self.generate_video_multi(
            left_audio_url=left_audio_url,
            right_audio_url=right_audio_url,
            image_url=image_url,
            order=order,
            resolution=resolution,
            seed=seed
        )

    def generate_video_multi(self, left_audio_url: str, right_audio_url: str, 
                            image_url: str, order: str = "meanwhile",
                            resolution: str = "480p", seed: int = -1) -> Optional[str]:
        """
        使用多个输入生成视频（多角色对话）
        
        Args:
            left_audio_url (str): 左侧说话人音频 URL
            right_audio_url (str): 右侧说话人音频 URL
            image_url (str): 背景图片 URL
            order (str): 音频顺序，可选 "meanwhile", "left_right", "right_left"
            resolution (str): 分辨率，如 "480p"
            seed (int): 随机种子，-1 表示随机
            
        Returns:
            str: 生成的视频 URL，失败时返回 None
        """
        print("Hello from WaveSpeedAI!")
        print(f"API_KEY: {self.api_key}")
        
        url = f"{self.base_url}/wavespeed-ai/infinitetalk/multi"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }
        payload = {
            "left_audio": left_audio_url,
            "right_audio": right_audio_url,
            "image": image_url,
            "order": order,
            "resolution": resolution,
            "seed": seed
        }
        
        print(f"请求参数: {payload}")
        
        try:
            begin = time.time()
            response = requests.post(url, headers=headers, data=json.dumps(payload))
            if response.status_code == 200:
                result = response.json()["data"]
                request_id = result["id"]
                print(f"Task submitted successfully. Request ID: {request_id}")
            else:
                print(f"Error: {response.status_code}, {response.text}")
                return None
            
            result_url = f"{self.base_url}/predictions/{request_id}/result"
            headers = {"Authorization": f"Bearer {self.api_key}"}
            
            # 轮询结果
            while True:
                response = requests.get(result_url, headers=headers)
                if response.status_code == 200:
                    result = response.json()["data"]
                    status = result["status"]
                    
                    if status == "completed":
                        end = time.time()
                        print(f"Task completed in {end - begin} seconds.")
                        video_url = result["outputs"][0]
                        print(f"Task completed. URL: {video_url}")
                        return video_url
                    elif status == "failed":
                        print(f"Task failed: {result.get('error')}")
                        return None
                    else:
                        print(f"Task still processing. Status: {status}")
                else:
                    print(f"Error: {response.status_code}, {response.text}")
                    return None
                
                time.sleep(0.1)
                
        except Exception as e:
            print(f"Error occurred: {str(e)}")
            return None
    
    def generate_video(self, audio_url: str, image_url: str, prompt: str = None,
                      resolution: str = "480p", seed: int = -1) -> Optional[str]:
        """
        生成视频
        
        Args:
            audio_url (str): 音频文件 URL
            image_url (str): 背景图片 URL
            prompt (str, optional): 正向提示词
            resolution (str): 分辨率，如 "480p"
            seed (int): 随机种子，-1 表示随机
            
        Returns:
            str: 生成的视频 URL，失败时返回 None
        """
        print("Hello from WaveSpeedAI!")
        print(f"API_KEY: {self.api_key}")
        
        url = f"{self.base_url}/wavespeed-ai/infinitetalk"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }
        payload = {
            "audio": audio_url,
            "image": image_url,
            "resolution": resolution,
            "seed": seed
        }
        
        # 添加可选参数
        if prompt:
            payload["prompt"] = prompt
        
        print(f"请求参数: {payload}")
        
        try:
            begin = time.time()
            response = requests.post(url, headers=headers, data=json.dumps(payload))
            if response.status_code == 200:
                result = response.json()["data"]
                request_id = result["id"]
                print(f"Task submitted successfully. Request ID: {request_id}")
            else:
                print(f"Error: {response.status_code}, {response.text}")
                return None
            
            result_url = f"{self.base_url}/predictions/{request_id}/result"
            headers = {"Authorization": f"Bearer {self.api_key}"}
            
            # 轮询结果
            while True:
                response = requests.get(result_url, headers=headers)
                if response.status_code == 200:
                    result = response.json()["data"]
                    status = result["status"]
                    
                    if status == "completed":
                        end = time.time()
                        print(f"Task completed in {end - begin} seconds.")
                        video_url = result["outputs"][0]
                        print(f"Task completed. URL: {video_url}")
                        return video_url
                    elif status == "failed":
                        print(f"Task failed: {result.get('error')}")
                        return None
                    else:
                        print(f"Task still processing. Status: {status}")
                else:
                    print(f"Error: {response.status_code}, {response.text}")
                    return None
                
                time.sleep(0.1)
                
        except Exception as e:
            print(f"Error occurred: {str(e)}")
            return None


def main():
    """主函数，用于测试 WaveSpeedAI 类"""
    try:
        wavespeed_ai = WaveSpeedAI()
        
        # 示例：上传文件
        # result = wavespeed_ai.upload_binary_file("/path/to/your/file")
        
        # 示例：生成视频
        # video_url = wavespeed_ai.generate_video()
        
        # 示例：生成多角色视频
        # video_url = wavespeed_ai.generate_video_multi()
        
    except ValueError as e:
        print(f"Initialization error: {e}")


if __name__ == "__main__":
    main()